package com.itaming.lycheeframework.security.token.provider;

import com.fasterxml.jackson.core.JsonGenerator;
import com.fasterxml.jackson.core.JsonParser;
import com.fasterxml.jackson.databind.*;
import com.fasterxml.jackson.databind.deser.std.UntypedObjectDeserializer;
import com.fasterxml.jackson.databind.module.SimpleModule;
import com.fasterxml.jackson.databind.ser.std.StdSerializer;
import com.itaming.lycheeframework.common.constant.FrameworkConstant;
import com.itaming.lycheeframework.security.exception.AuthenticationException;
import com.itaming.lycheeframework.security.exception.InvalidTokenException;
import com.itaming.lycheeframework.security.exception.TokenExpiredException;
import com.itaming.lycheeframework.security.token.TokenService;
import com.itaming.lycheeframework.support.utils.ArrayUtil;
import com.itaming.lycheeframework.support.utils.StringUtil;
import io.jsonwebtoken.Claims;
import io.jsonwebtoken.ExpiredJwtException;
import io.jsonwebtoken.Jwts;
import io.jsonwebtoken.jackson.io.JacksonDeserializer;
import io.jsonwebtoken.jackson.io.JacksonSerializer;
import io.jsonwebtoken.security.Keys;
import lombok.RequiredArgsConstructor;
import org.springframework.util.Assert;

import javax.crypto.SecretKey;
import java.io.IOException;
import java.util.Date;
import java.util.Map;

/**
 * JWT Token提供者
 *
 * @author A.Ming
 */
@RequiredArgsConstructor
public class JwtTokenProvider implements TokenProvider {

    /**
     * 发行者
     */
    private static final String ISSUER = FrameworkConstant.FRAMEWORK_NAME;

    /**
     * 类型
     */
    private static final String JWT = "JWT";

    /**
     * Claims中需要存储类型的字段
     */
    private static final String[] TYPED_CLAIMS_KEYS = {TokenService.PAYLOAD_KEY_ID};

    /**
     * ObjectMapper
     */
    private static final ObjectMapper OBJECT_MAPPER = createObjectMapper();

    /**
     * 序列化器
     */
    private static final JacksonSerializer<Map<String, ?>> CLAIMS_SERIALIZER = new JacksonSerializer<>(OBJECT_MAPPER);

    /**
     * 反序列化器
     */
    private static final JacksonDeserializer<Map<String, ?>> CLAIMS_DESERIALIZER = new JacksonDeserializer<>(OBJECT_MAPPER);

    /**
     * ObjectMapper模块名
     */
    private static final String OBJECT_MAPPER_MODULE_NAME = "lychee-jjwt-jackson";

    /**
     * 用来存储类型信息的键
     */
    private static final String TYPED_CLAIMS_TYPE_KEY = "@type";

    /**
     * 用来存储值信息的键
     */
    private static final String TYPED_CLAIMS_VALUE_KEY = "value";

    /**
     * 密钥
     */
    private final SecretKey secretKey;

    /**
     * 构造函数
     *
     * @param secretKey 密钥
     */
    public JwtTokenProvider(String secretKey) {
        Assert.isTrue(StringUtil.isNotEmpty(secretKey), "JWT secret key must not be empty");
        this.secretKey = Keys.hmacShaKeyFor(StringUtil.getBytes(secretKey));
    }

    /**
     * 创建Token
     *
     * @param payload   载荷
     * @param expiresIn 过期时间(秒）
     * @return Token
     */
    @Override
    public String createToken(Map<String, ?> payload, int expiresIn) {
        return Jwts.builder()
            .json(CLAIMS_SERIALIZER)
            .signWith(secretKey)
            .header()
            .type(JWT)
            .and()
            .issuer(ISSUER)
            .expiration(getExpireDate(expiresIn))
            .notBefore(new Date())
            .claims(payload)
            .compact();
    }

    /**
     * 解析
     *
     * @param token token
     * @return 载荷
     * @throws AuthenticationException 解析失败时抛出，如：Token过期、Token无效等
     */
    @Override
    public Map<String, Object> parseToken(String token) throws AuthenticationException {
        if (StringUtil.isEmpty(token)) {
            throw new AuthenticationException();
        }

        try {
            return parseClaims(token);
        } catch (ExpiredJwtException e) {
            throw new TokenExpiredException();
        } catch (Exception e) {
            throw new InvalidTokenException();
        }
    }

    /**
     * 获取过期时间
     *
     * @param expiresIn 过期时间(秒)
     * @return 过期时间
     */
    private Date getExpireDate(int expiresIn) {
        if (expiresIn <= 0) {
            throw new IllegalArgumentException("The token expiration time must be greater than 0");
        }

        return new Date(System.currentTimeMillis() + expiresIn * 1000L);
    }

    /**
     * 解析Claims
     *
     * @param token token
     * @return Claims
     */
    private Claims parseClaims(String token) {
        return Jwts.parser()
            .json(CLAIMS_DESERIALIZER)
            .verifyWith(secretKey)
            .build()
            .parseSignedClaims(token)
            .getPayload();
    }

    /**
     * 创建ObjectMapper
     *
     * @return ObjectMapper
     */
    private static ObjectMapper createObjectMapper() {
        ObjectMapper objectMapper = new ObjectMapper();

        SimpleModule module = new SimpleModule(OBJECT_MAPPER_MODULE_NAME);
        module.addSerializer(TypedClaimsSerializer.INSTANCE);
        module.addDeserializer(Object.class, TypedClaimsDeserializer.INSTANCE);

        return objectMapper
            .registerModule(module)
            .configure(JsonParser.Feature.STRICT_DUPLICATE_DETECTION, true)
            .configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
    }

    /**
     * 类型安全的Claims序列化器
     */
    private static class TypedClaimsSerializer extends StdSerializer<Map<String, ?>> {

        static final TypedClaimsSerializer INSTANCE = new TypedClaimsSerializer();

        private TypedClaimsSerializer() {
            super(Map.class, false);
        }

        @Override
        public void serialize(Map<String, ?> map, JsonGenerator generator, SerializerProvider provider) throws IOException {
            generator.writeStartObject();

            for (Map.Entry<String, ?> entry : map.entrySet()) {
                String key = entry.getKey();
                Object value = entry.getValue();

                // 针对特定字段保留类型信息
                if (ArrayUtil.contains(TYPED_CLAIMS_KEYS, key)) {
                    generator.writeFieldName(key);
                    generator.writeStartObject();
                    generator.writeStringField(TYPED_CLAIMS_TYPE_KEY, value.getClass().getName());
                    generator.writeObjectField(TYPED_CLAIMS_VALUE_KEY, value);
                    generator.writeEndObject();
                } else {
                    // 默认序列化其他字段
                    generator.writeObjectField(key, value);
                }
            }

            generator.writeEndObject();
        }

    }

    /**
     * 类型安全的Claims反序列化器
     */
    private static class TypedClaimsDeserializer extends UntypedObjectDeserializer {

        static final TypedClaimsDeserializer INSTANCE = new TypedClaimsDeserializer();

        private TypedClaimsDeserializer() {
            super((JavaType) null, null);
        }

        @Override
        public Object deserialize(JsonParser parser, DeserializationContext context) throws IOException {
            String key = parser.currentName();

            // 针对特定字段解析类型信息
            if (ArrayUtil.contains(TYPED_CLAIMS_KEYS, key)) {
                JsonNode node = parser.readValueAsTree();
                String type = node.get(TYPED_CLAIMS_TYPE_KEY).asText();
                JsonNode value = node.get(TYPED_CLAIMS_VALUE_KEY);
                try {
                    Class<?> clazz = Class.forName(type);
                    return parser.getCodec().treeToValue(value, clazz);
                } catch (Exception e) {
                    String message = String.format("Unable to deserialize key '%s' as type '%s'", key, type);
                    throw new IllegalArgumentException(message, e);
                }
            }

            // 默认反序列化其他字段
            return super.deserialize(parser, context);
        }

    }

}
